#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Kimi K2 AI助手 - 本地服务器
提供静态文件服务和CORS支持
"""

import os
import sys
import json
import asyncio
from pathlib import Path
from flask import Flask, send_from_directory, request, jsonify, Response
from flask_cors import CORS
import requests
from datetime import datetime
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 创建Flask应用
app = Flask(__name__)
CORS(app)  # 启用CORS支持

# 配置
CURRENT_DIR = Path(__file__).parent
STATIC_DIR = CURRENT_DIR
PORT = 8080
HOST = '127.0.0.1'

# Kimi API配置
KIMI_API_KEY = 'sk-vjna6ovnmsubah6fzxdczzme5lkpcf5zsy6dbkjczgb5uhog'  # 用户提供的API密钥
KIMI_BASE_URL = 'https://maas-api.lanyun.net/v1'
KIMI_MODEL = '/maas/kimi/Kimi-K2-Instruct'


@app.route('/')
def index():
    """主页"""
    return send_from_directory(STATIC_DIR, 'index.html')


@app.route('/<path:filename>')
def static_files(filename):
    """静态文件服务"""
    try:
        return send_from_directory(STATIC_DIR, filename)
    except FileNotFoundError:
        return jsonify({'error': 'File not found'}), 404


@app.route('/api/chat', methods=['POST'])
def chat_proxy():
    """聊天API代理（可选，用于解决CORS问题）"""
    try:
        data = request.get_json()
        logger.info(f"Received chat request: {data}")
        
        # 验证请求数据
        if not data or 'messages' not in data:
            logger.error("Invalid request data")
            return jsonify({'error': 'Invalid request data'}), 400
        
        # 构建请求
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {KIMI_API_KEY}'
        }
        
        payload = {
            'model': KIMI_MODEL,
            'messages': data['messages'],
            'stream': data.get('stream', True),
            'temperature': data.get('temperature', 0.7),
            'max_tokens': data.get('max_tokens', 2048)
        }
        
        # 发送请求到Kimi API
        logger.info(f"Sending request to Kimi API: {KIMI_BASE_URL}/chat/completions")
        response = requests.post(
            f'{KIMI_BASE_URL}/chat/completions',
            headers=headers,
            json=payload,
            stream=payload['stream']
        )
        
        logger.info(f"API response status: {response.status_code}")
        if response.status_code != 200:
            logger.error(f"API request failed: {response.status_code}, {response.text}")
            return jsonify({
                'error': f'API request failed: {response.status_code}',
                'message': response.text
            }), response.status_code
        
        # 如果是流式响应
        if payload['stream']:
            def generate():
                try:
                    for line in response.iter_lines():
                        if line:
                            line_text = line.decode('utf-8').strip()
                            logger.info(f"Raw line from API: {line_text}")
                            
                            # 只处理有效的SSE数据行
                            if line_text.startswith('data:'):
                                # 确保格式正确
                                if not line_text.startswith('data: '):
                                    line_text = line_text.replace('data:', 'data: ', 1)
                                yield line_text + '\n\n'
                            elif line_text and not line_text.startswith(':'):
                                # 跳过注释行和空行，只处理可能的数据行
                                # 检查是否是有效的JSON
                                try:
                                    import json
                                    json.loads(line_text)
                                    yield f"data: {line_text}\n\n"
                                except json.JSONDecodeError:
                                    # 不是有效的JSON，跳过
                                    logger.info(f"Skipping non-JSON line: {line_text}")
                                    pass
                except Exception as e:
                    logger.error(f"Stream processing error: {e}")
                    yield f"data: {{\"error\": \"{str(e)}\"}}\n\n"
            
            return Response(
                generate(),
                mimetype='text/event-stream',
                headers={
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                    'Access-Control-Allow-Origin': '*',
                    'Access-Control-Allow-Headers': 'Content-Type, Authorization',
                    'X-Accel-Buffering': 'no'
                }
            )
        else:
            return jsonify(response.json())
            
    except Exception as e:
        logger.error(f"Chat proxy error: {e}")
        return jsonify({'error': str(e)}), 500


@app.route('/api/health', methods=['GET'])
def health_check():
    """健康检查"""
    return jsonify({
        'status': 'healthy',
        'timestamp': datetime.now().isoformat(),
        'version': '1.0.0'
    })


@app.route('/api/config', methods=['GET'])
def get_config():
    """获取配置信息"""
    return jsonify({
        'model': KIMI_MODEL,
        'base_url': KIMI_BASE_URL,
        'features': {
            'stream': True,
            'markdown': True,
            'code_highlight': True
        }
    })


@app.errorhandler(404)
def not_found(error):
    """404错误处理"""
    return jsonify({'error': 'Not found'}), 404


@app.errorhandler(500)
def internal_error(error):
    """500错误处理"""
    logger.error(f"Internal server error: {error}")
    return jsonify({'error': 'Internal server error'}), 500


def check_dependencies():
    """检查依赖文件"""
    required_files = ['index.html', 'styles.css', 'script.js']
    missing_files = []
    
    for file in required_files:
        if not (STATIC_DIR / file).exists():
            missing_files.append(file)
    
    if missing_files:
        logger.warning(f"Missing files: {missing_files}")
        return False
    
    return True


def main():
    """主函数"""
    print("="*50)
    print("🚀 Kimi K2 AI助手 - 本地服务器")
    print("="*50)
    
    # 检查依赖
    if not check_dependencies():
        print("❌ 缺少必要的文件，请确保所有文件都在同一目录下")
        return
    
    print(f"📁 静态文件目录: {STATIC_DIR}")
    print(f"🌐 服务器地址: http://{HOST}:{PORT}")
    print(f"🤖 AI模型: {KIMI_MODEL}")
    print("-"*50)
    print("💡 使用说明:")
    print("   1. 在浏览器中打开上述地址")
    print("   2. 开始与AI助手对话")
    print("   3. 按 Ctrl+C 停止服务器")
    print("="*50)
    
    try:
        # 启动服务器
        app.run(
            host=HOST,
            port=PORT,
            debug=False,
            threaded=True
        )
    except KeyboardInterrupt:
        print("\n👋 服务器已停止")
    except Exception as e:
        logger.error(f"服务器启动失败: {e}")
        print(f"❌ 服务器启动失败: {e}")


if __name__ == '__main__':
    main()